//      mp.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include "stl-helper.h"
#include "mp.h"

void MP_SqLoss::print_result(const std::string& outputFile)
{
  FileAccessWrite(output, outputFile);
  for (int i = 0; i < sz(w_); i++) {
    output << w_[i] << ENDL;
  }
  output.close();
}

void MP_SqLoss::solve()
{
  data_.to_unit_norm();
  arma::vec r = data_.label();
  w_ = arma::zeros<arma::vec>(data_.dim());
  
  arma::vec old_r = r;
  selected_feature_.clear();
  for (int k = 0; ; k++) {
    int n_k = maxip(r);
    
    selected_feature_.insert(n_k);

    arma::mat X;
    build_data(X);
    solve_model(X, r);

    if (!debugFile_.empty()) {
      FileAccessAppend(debug, debugFile_);
      FileAccessAppendRow(debug, w_);
    }

    double new_norm = arma::norm(r, 2);

    MSG("iter " << k << " norm(r)=" << new_norm << " n=" << n_k );

    if (arma::norm(r-old_r,2) < 1e-6) break;
    old_r = r;
  }

}

int MP_SqLoss::maxip(const arma::vec& q)
{
  int max_i = -1;
  double max = -INFINITY;
  for (int i = 0; i < data_.dim(); i++) {
    double dot_q_x_i = arma::dot(data_.f(i), q) / arma::norm(data_.f(i), 2);
    //MSG(i << " " << dot_q_x_i);
    if (dot_q_x_i > max) {
      max_i = i;
      max = dot_q_x_i;
    }
  }
  return max_i;
}

void MP_SqLoss::build_data(arma::mat& X)
{
  X = arma::zeros<arma::mat>(data_.size(), selected_feature_.size()+1);
  for (uint i = 0; i < X.n_rows; i++) {
    int j = 0;
    BOOST_FOREACH(int n, selected_feature_) {
      X(i, j++) = data_.x(i, n);
    }
    X(i, j) = 1; // bias term
  }
}

void MP_SqLoss::solve_model(const arma::mat& X, arma::vec& r)
{
  arma::vec w = arma::solve(X, data_.label());
  r = data_.label() - X*w;
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  int j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w_[n] = w[j++];
  }
  w_[w_.n_elem-1] = w[w.n_elem-1]; // bias term
}

// ================ MP_LogLoss =========================
void MP_LogLoss::solve()
{
  if (vm_.count("C")) C_ = vm_["C"].as<double>();

  arma::vec r = 0.5*arma::ones<arma::vec>(data_.size()) % data_.label();
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  
  arma::vec old_r = r;
  selected_feature_.clear();
  if (!debugFile_.empty()) {
    FileAccessWrite(debug, debugFile_);
    debug.close();
  }
  for (int k = 0; sz(selected_feature_) < data_.dim(); k++) {
    int n_k = maxip(r);
    
    selected_feature_.insert(n_k);

    arma::mat X;
    build_data(X);
    solve_model(X, r);

    if (!debugFile_.empty()) {
      FileAccessAppend(debug, debugFile_);
      FileAccessAppendRow(debug, w_);
    }

    double new_norm = arma::norm(r, 2);

    MSG("iter " << k << " norm(r)=" << new_norm << " n=" << n_k );

    if (arma::norm(r-old_r,2) < 1e-6) break;
    old_r = r;
  }

}

arma::vec MP_LogLoss::classify(MatrixData& test)
{
  arma::vec w = arma::zeros<arma::vec>(sz(selected_feature_)+1);

  arma::mat X = arma::zeros<arma::mat>(sz(test), selected_feature_.size()+1);
  int j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    X.col(j) = test.f(n);
    w(j) = w_[n];
    j++;
  }
  for (int i = 0; i < sz(test); i++) X(i, j) = 1; // bias term
  w(j) = w_[sz(w_)-1];

  arma::vec f = X * w;
  for (int i = 0; i < sz(f); i++)
    f(i) = f(i) < 0 ? -1 : 1;
  return f;
}

void MP_LogLoss::print_result(const std::string& outputFile)
{
  FileAccessWrite(output, outputFile);
  for (int i = 0; i < sz(w_); i++) {
    output << w_[i] << ENDL;
  }
  output.close();
}

int MP_LogLoss::maxip(const arma::vec& q)
{
  int max_i = -1;
  double max = -INFINITY;
  for (int i = 0; i < data_.dim(); i++) {
    double dot_q_x_i = arma::dot(data_.f(i), q);
		
    //MSG(i << " " << dot_q_x_i);
    if (dot_q_x_i > max/* && selected_feature_.find(i) == selected_feature_.end()*/) {
      max_i = i;
      max = dot_q_x_i;
    }
  }
	
  return (max > 0) ? max_i : -1;
}

void MP_LogLoss::build_data(arma::mat& X)
{
  X = arma::zeros<arma::mat>(data_.size(), selected_feature_.size()+1);
  for (uint i = 0; i < X.n_rows; i++) {
    int j = 0;
    BOOST_FOREACH(int n, selected_feature_) {
      X(i, j++) = data_.x(i, n);
    }
    X(i, j) = 1; // bias term
  }
}

void MP_LogLoss::fval_gradient(const arma::mat& X, const arma::vec& w, 
			       arma::vec& f, arma::vec& e, arma::vec& p, double& f_val)
{
  f = X*w;
  e =  arma::exp(-data_.label() % f);
  p = e+1;
  f_val = 0;
  for (int i = 0; i < sz(p); i++) {
    p(i) = e(i) / p(i) * data_.y(i);
    f_val += log(e(i)+1);
  }
  f_val = C_/data_.size() * f_val + 0.5*arma::dot(w,w);
}

void MP_LogLoss::solve_model(const arma::mat& X, arma::vec& p)
{
  arma::vec w = arma::zeros<arma::vec>(selected_feature_.size()+1), f, e;
  // copy the current weight vector to initialize
  int j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w[j++] = w_[n];
  }
  w[w.n_elem-1] = w_[w_.n_elem-1]; // bias term

  double f_value;
  fval_gradient(X, w, f, e , p, f_value);

  for (int iter = 0; iter < 10000; iter++) {
    arma::vec dLdw = C_/data_.size()*arma::trans(X)*p;
    //MSG("dLdw = " << dLdw);
    for (int i = 0; i < sz(w)-1; i++) dLdw(i) -= w(i);

    // line search
    double tol = 1e-6, alpha = 1;
    while (alpha >= tol) {
      arma::vec wa = w + alpha * dLdw, fa, ea, pa;
      double f_a;
      fval_gradient(X, wa, fa, ea, pa, f_a);

      if (f_a < f_value) {
	w = wa;
	f = fa;
	e = ea;
	p = pa;
	f_value = f_a;

	break;
      }
      alpha *= 0.5;
    }
    //MSG(iter << " " << arma::sum(arma::log(e+1)));
    //MSG("w = " << w);
    // step too small
    if (alpha < tol) {
      MSG("step too small alpha=" << alpha);
      break;
    }
  }
  
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w_[n] = w[j++];
  }
  w_[w_.n_elem-1] = w[w.n_elem-1]; // bias term
}

// ================ MP_SqHingeLoss =========================
void MP_SqHingeLoss::solve()
{
  C_ = vm_["C"].as<double>();

  arma::vec r = 0.5*arma::ones<arma::vec>(data_.size()) % data_.label();
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  
  arma::vec old_r = r;
  selected_feature_.clear();
  if (!debugFile_.empty()) {
    FileAccessWrite(debug, debugFile_);
    debug.close();
  }
  for (int k = 0; sz(selected_feature_) < data_.dim(); k++) {
    int n_k = maxip(r);
    
    selected_feature_.insert(n_k);

    arma::mat X;
    build_data(X);
    solve_model(X, r);

    if (!debugFile_.empty()) {
      FileAccessAppend(debug, debugFile_);
      FileAccessAppendRow(debug, w_);
    }

    double new_norm = arma::norm(r, 2);

    MSG("iter " << k << " norm(r)=" << new_norm << " n=" << n_k );

    if (arma::norm(r-old_r,2) < 1e-6 || norm(r, 2) < 1e-6) break;
    old_r = r;
  }

}

void MP_SqHingeLoss::fval_gradient(const arma::mat& X, const arma::vec& w, arma::vec& f, arma::vec& e, arma::vec& p, double& f_val)
{
  f = X*w;
  e = 1 - data_.label()%f;
  for (int i = 0; i < sz(e); i++)
    if (e(i) < 0) e(i) = 0;
  p = data_.label() % e;

  f_val = C_/sz(data_)*arma::dot(e,e)+0.5*arma::dot(w,w);
}

void MP_SqHingeLoss::solve_model(const arma::mat& X, arma::vec& p)
{
  arma::vec w = arma::zeros<arma::vec>(selected_feature_.size()+1), f, e;

  double f_value;
  fval_gradient(X,w,f,e,p,f_value);
  for (int iter = 0; iter < 10000; iter++) {
    arma::vec dLdw = C_/data_.size()*arma::trans(X)*p;
    //MSG("dLdw = " << dLdw);
    for (int i = 0; i < sz(w)-1; i++) dLdw(i) -= w(i);

    // line search
    double tol = 1e-16, alpha = 1;
    while (alpha >= tol) {
      arma::vec wa = w + alpha * dLdw, fa, ea, pa;
      double f_a;
      fval_gradient(X,wa,fa,ea,pa,f_a);
      //MSG("wa=" << wa << "   " << alpha << " f_a = " << f_a);
      if (f_a < f_value) {
	w = wa;
	f = fa;
	e = ea;
	p = pa;
	f_value = f_a;
	break;
      }
      alpha *= 0.5;
    }
    //MSG("descent " << iter << " " << f_value);
    // step too small
    if (alpha < tol) {
      MSG("iter=" << iter << " step too small alpha=" << alpha);
      break;
    }
  }
  
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  int j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w_[n] = w[j++];
  }
  w_[w_.n_elem-1] = w[w.n_elem-1]; // bias term
}

// ================ MP_SmoothHingeLoss =========================
void MP_SmoothHingeLoss::solve()
{
  if (vm_.count("C")) C_ = vm_["C"].as<double>();
  if (vm_.count("mu")) mu_ = vm_["mu"].as<double>();

  arma::vec r = arma::ones<arma::vec>(data_.size());
  if (1.0 / mu_ < 1) r *= 1.0 / mu_;
  r %= data_.label();	
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  
  arma::vec old_r = r;
  selected_feature_.clear();
	
  if (!debugFile_.empty()) {
    FileAccessWrite(debug, debugFile_);
    debug.close();
  }
  for (int k = 0; sz(selected_feature_) < data_.dim(); k++) {
    int n_k = maxip(r);
		if (n_k < 0) {
			MSG("no more good feature, terminate.");
			break;
		}
    
    selected_feature_.insert(n_k);

    arma::mat X;
    build_data(X);
    solve_model(X, r);

    if (!debugFile_.empty()) {
      FileAccessAppend(debug, debugFile_);
      FileAccessAppendRow(debug, w_);
    }

    double new_norm = arma::norm(r, 2);

    MSG("iter " << k << " norm(r)=" << new_norm << " n=" << n_k );

    if (arma::norm(r-old_r,2) < 1e-6 /*|| norm(r, 2) < 1e-6*/) break;
    old_r = r;
  }

}

void MP_SmoothHingeLoss::solve_model(const arma::mat& X, arma::vec& p)
{
  arma::vec w = arma::zeros<arma::vec>(selected_feature_.size()+1),
    f, e;
  double f_value;

  // copy the current weight vector to initialize
  int j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w[j++] = w_[n];
  }
  w[w.n_elem-1] = w_[w_.n_elem-1]; // bias term

  fval_gradient(X, w, f, e, p, f_value);
	
  for (int iter = 0; iter < 10000; iter++) {
    arma::vec dLdw = C_ / data_.size() * arma::trans(X)*p;
    for (int i = 0; i < sz(w)-1; i++) dLdw(i) -= w(i);
    //MSG(iter << " w= " << w << "f_value = " << f_value);

    // line search
    double tol = 1e-16, alpha = 1;
    while (alpha >= tol) {
      arma::vec wa = w + alpha * dLdw, fa, ea, pa;
      double f_a;
      fval_gradient(X, wa, fa, ea, pa, f_a);

      //MSG("wa=" << wa << "   " << alpha << " f_a = " << f_a);
      if (f_a < f_value) {
	w = wa;
	f = fa;
	e = ea;
	p = pa;
	f_value = f_a;
	break;
      }
      alpha *= 0.5;
    }
    //MSG("descent " << iter << " " << f_value << " " << 0.5*arma::dot(w,w) << " p=" << arma::trans(p) << " f=" << arma::trans(f));
    // step too small
    if (alpha < tol) {
      MSG("iter=" << iter << " step too small alpha=" << alpha);
      break;
    }
  }
  
  w_ = arma::zeros<arma::vec>(data_.dim()+1);
  j = 0;
  BOOST_FOREACH(int n, selected_feature_) {
    w_[n] = w[j++];
  }
  w_[w_.n_elem-1] = w[w.n_elem-1]; // bias term
}

void MP_SmoothHingeLoss::fval_gradient(const arma::mat& X, const arma::vec& w, arma::vec& f, arma::vec& e, arma::vec& p, double& f_val)
{
  f = X*w;
  p = e = (1 - data_.label()%f) / mu_;
  f_val = 0;
  for (int i = 0; i < sz(e); i++) {
    if (e(i) < 0) p(i) = 0;
    else if (e(i) < 1) {
      p(i) = e(i);
      f_val += e(i)*e(i)*mu_/2;
    }
    else {
      p(i) = 1;
      f_val += e(i)*mu_-mu_/2;
    }
  }
  f_val = f_val * C_ / data_.size() + 0.5 * arma::dot(w, w);
  p %= data_.label();
}

void MP_SmoothHingeLoss::solve_all_feature()
{
  if (vm_.count("C")) C_ = vm_["C"].as<double>();
  if (vm_.count("mu")) mu_ = vm_["mu"].as<double>();

  w_ = arma::zeros<arma::vec>(data_.dim()+1);
	selected_feature_.clear();
	for (int i = 0; i < data_.dim(); i++) selected_feature_.insert(i);
	
	arma::mat X;
	arma::vec r;
	build_data(X);
	solve_model(X, r);
}
